import logging
from sklearn.model_selection import train_test_split
from tqdm import tqdm
from sklearn.datasets import load_digits
from collections import Counter
from pinecone import Pinecone, ServerlessSpec

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

# 初始化 Pinecone 客户端
pinecone = Pinecone(api_key="120457e3-cd8e-4d78-8208-6b016c02bf75")

index_name = "mnist-index"

# 检查索引是否存在，如果存在则删除
existing_indexes = pinecone.list_indexes()
if any(index['name'] == index_name for index in existing_indexes):
    pinecone.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已存在，已删除。")

print(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,
    metric="euclidean",
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

index = pinecone.Index(index_name)

# 加载 MNIST 数据集
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 将数据转换为 Pinecone 需要的格式
vectors = []
for i in tqdm(range(len(X_train)), desc="准备训练数据"):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))

# 批量插入数据
batch_size = 1000
with tqdm(total=len(vectors), desc="上传数据") as pbar:
    for i in range(0, len(vectors), batch_size):
        batch = vectors[i:i + batch_size]
        index.upsert(batch)
        pbar.update(batch_size)
logging.info(f"成功创建索引，并上传了{len(vectors)}条数据")

# 测试准确率
correct_predictions = 0
total_predictions = 0

logging.info("开始测试准确率...")
for i, d in enumerate(X_test):
    results = index.query(vector=d.tolist(), top_k=11, include_metadata=True)
    labels = [int(match['metadata']['label']) for match in results['matches']]
    if labels:  # 检查是否有标签数据
        prediction = Counter(labels).most_common(1)[0][0]
        if prediction == y_test[i]:
            correct_predictions += 1
    total_predictions += 1
    if i % 100 == 0:
        logging.info(f"已测试 {i} 条数据...")
logging.info(f"当 k=11 时，使用 Pinecone 的准确率为: {correct_predictions / total_predictions:.2f}")
